import time
import cv2
import numpy as np
from PIL import Image

from yolo import YOLO

if __name__ == "__main__":
    yolo = YOLO()

    mode = "dir_predict"
    crop = False
    count = False
    video_path = 0
    video_save_path = ""
    video_fps = 25.0
    dir_origin_path = r"I:\SMTData\Data_旗开得\3D\IC_NG"
    dir_save_path = r"I:\SMTData\Data_旗开得\3D\IC_NG"
    heatmap_save_path = "model_data/heatmap_vision.png"
    simplify = True
    onnx_save_path = "model_data/models.onnx"
    result_file = r"I:\SMTData\Data_旗开得\3D\IC_NG\result_NG.txt"

    class_names = ['FOREIGN', 'SOLDER', 'BRIDGE', 'XIZHU', 'SHIFT']
    OK_result = "OK"

    if mode == "dir_predict":
        import os
        from tqdm import tqdm

        img_names = os.listdir(dir_origin_path)
        with open(result_file, "a+") as f:
            for img_name in tqdm(img_names):
                if img_name.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm')):
                    image_path = os.path.join(dir_origin_path, img_name)
                    image = Image.open(image_path)
                    r_image = yolo.detect_image(image)
                    r_image.save(os.path.join(dir_save_path, img_name.replace(".jpg", ".png")), quality=95, subsampling=0)
                    res = yolo.detect_imageOne(image, crop=crop, count=count)
                    if res is not None:

                        top_label, top_conf, top_boxes = res
                        num_boxes = len(top_boxes)
                        ng_flag = False  # 重置 ng_flag
                        if num_boxes > 0:
                            for i, c in list(enumerate(top_label)):
                                predicted_class = class_names[int(c)]
                                box = top_boxes[i]
                                score = top_conf[i]
                                top, left, bottom, right = box
                                top = max(0, np.floor(top).astype('int32'))
                                left = max(0, np.floor(left).astype('int32'))
                                bottom = min(image.size[1], np.floor(bottom).astype('int32'))
                                right = min(image.size[0], np.floor(right).astype('int32'))
                                imgOne = cv2.imread(image_path)
                                if predicted_class == 'XIZHU':
                                    cv2.rectangle(imgOne, (left, top), (right, bottom), (0, 0, 255), 2)
                                    f.write(f"{img_name,predicted_class,score},NG\n")
                                elif predicted_class == 'SHIFT' and score >= 0.5:
                                    cv2.rectangle(imgOne, (left, top), (right, bottom), (255, 0, 0), 2)
                                    # f.write(f"{img_name, predicted_class, score},NG\n")
                                    f.write(f"{img_name},NG\n")
                                elif predicted_class == 'FOREIGN' and score >= 0.1:
                                    cv2.rectangle(imgOne, (left, top), (right, bottom), (255, 0, 0), 2)
                                    f.write(f"{img_name},NG\n")
                                elif predicted_class == 'BRIDGE' and score >= 0.1:
                                    cv2.rectangle(imgOne, (left, top), (right, bottom), (255, 0, 0), 2)
                                    f.write(f"{img_name},NG\n")
                                elif predicted_class == 'SOLDER' and score >= 0.1:
                                    cv2.rectangle(imgOne, (left, top), (right, bottom), (255, 0, 0), 2)
                                    f.write(f"{img_name},NG\n")
                        else:
                            f.write(f"{img_name},OK\n")
                    else:
                        # 如果没有检测到物体
                        f.write(f"{img_name},OK\n")